import sox

import subprocess
from pathlib import Path
import librosa
from scipy.io import wavfile
import numpy as np
import torch
import csv
import whisper
import os
import shutil
import sys
import time
import datetime

dataset_id = "/root/data/output_training_data/"
pretrain_work_dir = "/root/data/pretrain_work_dir/"
wavs_work_dir = "/root/data/test_wavs/"
audio_file = "/root/data/aud.wav"

def clear_or_create_directory(dir_path):
    if os.path.exists(dir_path):
        for filename in os.listdir(dir_path):
            file_path = os.path.join(dir_path, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f'Failed to delete {file_path}. Reason: {e}')
    else:
        try:
            os.makedirs(dir_path)
        except OSError as e:
            print(f'Failed to create directory {dir_path}. Reason: {e}')

def split_long_audio(model, filepaths, save_dir="data_dir", out_sr=44100):
    if isinstance(filepaths, str):
        filepaths = [filepaths]

    for file_idx, filepath in enumerate(filepaths):

        save_path = Path(save_dir)
        save_path.mkdir(exist_ok=True, parents=True)

        print(f"Transcribing file {file_idx}: '{filepath}' to segments...")
        result = model.transcribe(filepath, word_timestamps=True, task="transcribe", beam_size=5, best_of=5)
        segments = result['segments']

        wav, sr = librosa.load(filepath, sr=None, offset=0, duration=None, mono=True)
        wav, _ = librosa.effects.trim(wav, top_db=20)
        peak = np.abs(wav).max()
        if peak > 1.0:
            wav = 0.98 * wav / peak
        wav2 = librosa.resample(wav, orig_sr=sr, target_sr=out_sr)
        wav2 /= max(wav2.max(), -wav2.min())

        for i, seg in enumerate(segments):
            start_time = seg['start']
            end_time = seg['end']
            wav_seg = wav2[int(start_time * out_sr):int(end_time * out_sr)]
            wav_seg_name = f"{file_idx}_{i}.wav"
            out_fpath = save_path / wav_seg_name
            wavfile.write(out_fpath, rate=out_sr, data=(wav_seg * np.iinfo(np.int16).max).astype(np.int16))
            
whisper_size = "medium"
whisper_model = whisper.load_model(whisper_size)

from modelscope.tools import run_auto_label

import os
from modelscope.models.audio.tts import SambertHifigan
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from modelscope.metainfo import Trainers
from modelscope.trainers import build_trainer
from modelscope.utils.audio.audio_utils import TtsTrainType

pretrained_model_id = 'damo/speech_personal_sambert-hifigan_nsf_tts_zh-cn_pretrain_16k'

def auto_label(audio):
    try:
        split_long_audio(whisper_model, audio, wavs_work_dir)
        input_wav = wavs_work_dir
        output_data = dataset_id
        ret, report = run_auto_label(input_wav=input_wav, work_dir=output_data, resource_revision="v1.0.7")

    except Exception:
        pass
    return "标注成功"

def train(a):
    try:

        train_info = {
            TtsTrainType.TRAIN_TYPE_SAMBERT: {  # 配置训练AM（sambert）模型
                'train_steps': 52,               # 训练多少个step
                'save_interval_steps': 50,       # 每训练多少个step保存一次checkpoint
                'log_interval': 10               # 每训练多少个step打印一次训练日志
            }
        }

        # 配置训练参数，指定数据集，临时工作目录和train_info
        kwargs = dict(
            model=pretrained_model_id,                  # 指定要finetune的模型
            model_revision = "v1.0.6",
            work_dir=pretrain_work_dir,                 # 指定临时工作目录
            train_dataset=dataset_id,                   # 指定数据集id
            train_type=train_info                       # 指定要训练类型及参数
        )

        trainer = build_trainer(Trainers.speech_kantts_trainer,
                                default_args=kwargs)

        trainer.train()

    except Exception:
        pass

    return "训练完成"


import random

def infer(text,filename):

    model_dir = pretrain_work_dir
    
    custom_infer_abs = {
        'voice_name':
        'F7',
        'am_ckpt':
        os.path.join(model_dir, 'tmp_am', 'ckpt'),
        'am_config':
        os.path.join(model_dir, 'tmp_am', 'config.yaml'),
        'voc_ckpt':
        os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan', 'ckpt'),
        'voc_config':
        os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan',
                'config.yaml'),
        'audio_config':
        os.path.join(model_dir, 'data', 'audio_config.yaml'),
        'se_file':
        os.path.join(model_dir, 'data', 'se', 'se.npy')
    }
    kwargs = {'custom_ckpt': custom_infer_abs}

    model_id = SambertHifigan(os.path.join(model_dir, "orig_model"), **kwargs)

    inference = pipeline(task=Tasks.text_to_speech, model=model_id)
    output = inference(input=text)

    if not filename:
        filename = str(random.randint(1, 1000000000000)) + "myfile.wav"

    with open(filename, mode='bx') as f:
        f.write(output["output_wav"])
    return filename

def launch_gradio(port=6006):
    import gradio as gr

    app = gr.Blocks()

    with app:
        gr.Markdown("# <center>🥳🎶🎡 - Sambert中文声音克隆</center>")
        gr.Markdown("## <center>🌟 - 训练3分钟，推理5秒钟，中英自然发音 </center>")
        gr.Markdown("### <center>🌊 - 更多精彩应用，敬请关注[滔滔AI](http://www.talktalkai.com)；滔滔AI，为爱滔滔！💕</center>")

        with gr.Row():
            inp1 = gr.Audio(type="filepath", label="请上传一段音频")
            out1 = gr.Textbox(label="标注情况", lines=1, interactive=False)

            out2 = gr.Textbox(label="训练情况", lines=1, interactive=False)
            inp2 = gr.Textbox(label="文本", lines=3)
            out3 = gr.Audio(type="filepath", label="合成的音频")
            btn1 = gr.Button("1.标注数据")
            btn2 = gr.Button("2.开始训练")
            btn3 = gr.Button("3.一键推理", variant="primary")

            btn1.click(auto_label, inp1, out1)
            btn2.click(train, out1, out2)
            btn3.click(infer, inp2, out3)

        gr.Markdown("### <center>注意❗：请不要生成会对个人以及组织造成侵害的内容，此程序仅供科研、学习及个人娱乐使用。</center>")
        gr.HTML('''
            <div class="footer">
                        <p>🌊🏞️🎶 - 江水东流急，滔滔无尽声。 明·顾璘
                        </p>
            </div>
        ''')
    app.launch(show_error=True, server_name="0.0.0.0", server_port=port)

def train_model():
    clear_or_create_directory(dataset_id)
    clear_or_create_directory(pretrain_work_dir)
    clear_or_create_directory(pretrain_work_dir)
    ret = auto_label(audio_file)
    print(ret)
    ret = train(audio_file)
    print(ret)
    # with open('/root/data/aud.wav', 'rb') as audio_file:
    #     auto_label(audio_file)
    #     train(audio_file)

def generate_audio(txt):
    current_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    output_filename = f"/root/data/result-{current_time}.wav"
    infer(txt,output_filename)
    print(f"Generating audio file: {output_filename}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python app.py gui/train/gen [text]")
    else:
        arg = sys.argv[1]
        if arg == 'gui':
            launch_gradio()
        elif arg == 'train':
            train_model()
        elif arg == 'gen':
            generate_audio(sys.argv[2])
        else:
            print("Usage: python app.py gui/train/gen [text]")